#  begin by clearing memory
rm(list=ls())
#  Load libraries
library(cmdstanr)
library(posterior)
library(readxl)
# Set directory
setwd("/Users/mwatson/Dropbox/TVExtreme/ReplicationFiles/Stan/WeatherDamages")

# --- Run for raw damages then for normalized damages 
# Specify model input
stan_dir <- "/Users/mwatson/Dropbox/TVExtreme/ReplicationFiles/Stan/Stan_Code"
application_dir = "/Users/mwatson/Dropbox/TVExtreme/ReplicationFiles/Stan/WeatherDamages"
model_str <- "Exceedance_GEV_4RW_Model"

# ==== Damages Model ... un-normalized ====
application_str = "Disasters"
data_str <- "Disasters"
data_file <- file.path(application_dir, paste0(data_str,"_y.xlsx"))
y <- read_excel(data_file,col_names = FALSE)
data_file <- file.path(application_dir, paste0(data_str,"_tau.xlsx"))
tau <- read_excel(data_file,col_names = FALSE)
data_file <- file.path(application_dir, paste0(data_str,"_nobs.xlsx"))
nobs <- read_excel(data_file,col_names = FALSE)

# # ------ Damages Model ... normalized ------
# application_str = "Disasters_Normalized"
# data_str <- "Disasters"
# data_file <- file.path(application_dir, paste0(data_str,"_y_norm.xlsx"))
# y <- read_excel(data_file,col_names = FALSE)
# data_file <- file.path(application_dir, paste0(data_str,"_tau_norm.xlsx"))
# tau <- read_excel(data_file,col_names = FALSE)
# data_file <- file.path(application_dir, paste0(data_str,"_nobs.xlsx"))
# nobs <- read_excel(data_file,col_names = FALSE)

# Process data and get dimensions
# Convert to matrix
y <- as.matrix(y)
y <- t(y)
# Convert nobs to a vector
nobs <- as.matrix(nobs) 
nobs <- c(nobs)
# Convert tau to a vector
tau <- as.matrix(tau) 
tau <- c(tau)
#  Get dimensions of cs_data
T <- dim(y)[1]

# Parameter values
#  xi bounds
x_min <- -0.9
x_max <- 2.1

# Scale gamma parameters
sg_xi <- 1.0
sg_alpha <- 0.5
sg_s <- 0.5

# Create list of data inputs for Stan
data_list <- list(
    T = T,
    xi_min = x_min,
    xi_max = x_max,
    sg_xi = sg_xi,
    sg_alpha = sg_alpha,
    sg_s = sg_s,
    tau = tau,
    nobs = nobs,
    y = t(y)
)

#  Inittial Values of parameters and seed value
trans_xi_level_init <- 0.30
ln_s_level_init <- -0.15
m_level_init <- 0.45
g_xi_init <- 0.2
g_alpha_init <- 0.2
g_s_init <- 0.2
g_m_init <- 0.2


init_list_1 <- list(
    trans_xi_level = trans_xi_level_init,
    ln_s_level = ln_s_level_init,
    m_level = m_level_init,
    g_xi = g_xi_init,
    g_alpha = g_alpha_init,
    g_s = g_s_init,
    g_m = g_m_init
)
init_list_2 <- list(
    trans_xi_level = trans_xi_level_init + 0.1,
    ln_s_level = ln_s_level_init + 0.1,
    m_level = m_level_init + 0.1,
    g_xi = g_xi_init,
    g_alpha = g_alpha_init,
    g_s = g_s_init,
    g_m = g_m_init
)
init_list_3 <- list(
    trans_xi_level = trans_xi_level_init - 0.1,
    ln_s_level = ln_s_level_init - 0.1,
    m_level = m_level_init - 0.1,
    g_xi = g_xi_init,
    g_alpha = g_alpha_init,
    g_s = g_s_init,
    g_m = g_m_init
)
init_list_4 <- list(
    trans_xi_level = trans_xi_level_init + 0.0,
    ln_s_level = ln_s_level_init + 0.1,
    m_level = m_level_init + 0.1,
    g_xi = g_xi_init,
    g_alpha = g_alpha_init,
    g_s = g_s_init,
    g_m = g_m_init
)
seed_value <- 87123

#  Model file
model_file <- file.path(stan_dir, paste0(model_str, ".stan"))

# Output Files
draws_file <- file.path(application_dir, paste0("CSV/",application_str,".",model_str,".Draws."))
output_file <- file.path(application_dir, paste0("Output/",application_str,".",model_str))

# Compile model
mod <- cmdstan_model(model_file)

# MCMC Draws model
fit <- mod$sample(
    data = data_list,
    iter_warmup = 5000,
    iter_sampling = 10000,
    init=list(
        init_list_1, # chain 1
        init_list_2, # chain 2
        init_list_3, # chain 3
        init_list_4  # chain 4
    ),
    refresh = 5000,
    chains = 4,
    parallel_chains = 4,
    seed = seed_value
) 
# fit diagnostics
sink(paste0(output_file,".diagonstics.txt"))
fit$diagnostic_summary()
#  fit summary
sink(paste0(output_file,".summary.txt"))
fit$print(max_rows = 10000,variables = c("gamma_xi","gamma_alpha","gamma_s","gamma_m","m_level","ln_s_level","xi","alpha","mu","sigma"))
sink(file = NULL)

# Save Draws of key parameters
draws_arr <- fit$draws()

tmp <- subset_draws(draws_arr, c("gamma_xi"))
gamma_xi_draws <- merge_chains(tmp)
write.csv(gamma_xi_draws, paste0(draws_file, "gamma_xi.csv"))
#
tmp <- subset_draws(draws_arr, c("gamma_alpha"))
gamma_alpha_draws <- merge_chains(tmp)
write.csv(gamma_alpha_draws, paste0(draws_file, "gamma_alpha.csv"))
#
tmp <- subset_draws(draws_arr, c("gamma_s"))
gamma_s_draws <- merge_chains(tmp)
write.csv(gamma_s_draws, paste0(draws_file, "gamma_s.csv"))
#
tmp <- subset_draws(draws_arr, c("gamma_m"))
gamma_m_draws <- merge_chains(tmp)
write.csv(gamma_m_draws, paste0(draws_file, "gamma_m.csv"))
#
tmp <- subset_draws(draws_arr, c("xi"))
xi_draws <- merge_chains(tmp)
write.csv(xi_draws, paste0(draws_file, "xi.csv"))
#
tmp <- subset_draws(draws_arr, c("alpha"))
alpha_draws <- merge_chains(tmp)
write.csv(alpha_draws, paste0(draws_file, "alpha.csv"))
#
tmp <- subset_draws(draws_arr, c("sigma"))
sigma_draws <- merge_chains(tmp)
write.csv(sigma_draws, paste0(draws_file, "sigma.csv"))
#
tmp <- subset_draws(draws_arr, c("mu"))
mu_draws <- merge_chains(tmp)
write.csv(mu_draws, paste0(draws_file, "mu.csv"))
